import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
from adjustText import adjust_text


def plot_dim(res, dim, idx, label_list, title, label):
    # idx_list = []
    # for i, idxi in enumerate(idx):
    #     if i != dim:
    #         idx_list.append([idxi])
    #     else:
    #         idx_list.append(range(res.shape[dim]))
    dim_list = list(range(len(res.shape)))
    dim_list.remove(dim)
    data = cal_mean(res, dim_list)
    fig, ax = plt.subplots()
    ax.scatter(label_list[dim], data)
    title += "-" + label
    ax.set_title(title)
    ax.set_xlabel(label)
    adjust_text([ax.text(i, j, f"{j:.4f}") for i, j in zip(label_list[dim], data)])
    fig.savefig(os.path.join("./res", title + ".png"))


def cal_mean(arr, dim_list):
    res = arr
    dim_list = sorted(dim_list, reverse=True)
    num = res != 0
    for dim in dim_list:
        num = num.sum(axis=dim)
        res = res.sum(axis=dim)
    res = res / num
    res[res == np.nan] = 0
    res[res == np.inf] = 0
    return res


def rebuild_idx(n, shape):
    idx = []
    for m in shape:
        idx.append(n % m)
        n //= m
    return idx


def get_label(idx, label_list, label_name):
    label = {}
    for i, idxi in enumerate(idx):
        label[label_name[i]] = label_list[i][idxi]
    return label


def get_one(pth):
    res = pd.read_csv(pth, sep="\t")
    return res["recall@5"].max()


def get_baseline():
    base_pth = "./trained_model/Douban/KG_free"
    embed_dim_list = [8, 16, 32, 64]
    lr_list = [0.1, 0.01, 0.001]
    res_table = np.zeros((len(embed_dim_list), len(lr_list)))
    for i0, dim in enumerate(embed_dim_list):
        for i1, lr in enumerate(lr_list):
            pth = os.path.join(base_pth, f"dim{dim}_lr{lr}_l20.0001", "metrics0.tsv")
            try:
                res = get_one(pth)
            except:
                res = 0
            res_table[i0, i1] = res
    return res_table


def analyze_base(res):
    embed_dim_list = [8, 16, 32, 64]
    lr_list = [0.1, 0.01, 0.001]

    idx = np.unravel_index(np.argmax(res), res.shape)

    label_list = (embed_dim_list, lr_list)
    label_name = ("embed_dim", "lr")
    label = get_label(idx, label_list, label_name)
    for dim in range(len(label_list)):
        plot_dim(res, dim, idx, label_list, "base", label_name[dim])
    return res[tuple(idx)], label


def get_embedding():
    base_pth = "./trained_model/Douban/Embedding_based"
    embed_dim_list = [8, 16, 32, 64]
    lr_list = [0.1, 0.01, 0.001]
    embedding_type_list = ["TransE", "TransR"]
    res_table = np.zeros((len(embed_dim_list), len(lr_list), len(embedding_type_list)))
    for i0, dim in enumerate(embed_dim_list):
        for i1, lr in enumerate(lr_list):
            for i2, em_type in enumerate(embedding_type_list):
                pth = os.path.join(base_pth, f"dim{dim}_lr{lr}_l20.0001_{em_type}", "metrics0.tsv")
                try:
                    res = get_one(pth)
                except:
                    res = 0
                res_table[i0, i1, i2] = res
    return res_table


def analyze_embedding(res):
    embed_dim_list = [8, 16, 32, 64]
    lr_list = [0.1, 0.01, 0.001]
    embedding_type_list = ["TransE", "TransR"]

    idx = np.unravel_index(np.argmax(res), res.shape)

    label_list = (embed_dim_list, lr_list, embedding_type_list)
    label_name = ("embed_dim", "lr", "embedding_type")
    label = get_label(idx, label_list, label_name)
    for dim in range(len(label_list)):
        plot_dim(res, dim, idx, label_list, "embedding", label_name[dim])
    return res[tuple(idx)], label


def get_gnn():
    base_pth = "./trained_model/Douban/GNN_based"
    embed_dim_list = [16, 32, 64]
    lr_list = [0.1, 0.01, 0.001]
    embedding_type_list = ["TransE", "TransR"]
    laplacian_type_list = ["symmetric", "random-walk"]
    aggregation_type_list = ["gcn", "graphsage", "lightgcn"]
    layers_list = [2, 4, 8]
    res_table = np.zeros(
        (
            len(embed_dim_list),
            len(lr_list),
            len(embedding_type_list),
            len(laplacian_type_list),
            len(aggregation_type_list),
            len(layers_list),
        )
    )
    for i0, dim in enumerate(embed_dim_list):
        for i1, lr in enumerate(lr_list):
            for i2, em_type in enumerate(embedding_type_list):
                for i3, lap_type in enumerate(laplacian_type_list):
                    for i4, agg_type in enumerate(aggregation_type_list):
                        for i5, nlayers in enumerate(layers_list):
                            pth = os.path.join(
                                base_pth,
                                f"nlayer{nlayers}_dim{dim}_lr{lr}_{lap_type}_{em_type}_{agg_type}",
                                "metrics0.tsv",
                            )
                            try:
                                res = get_one(pth)
                            except:
                                res = 0
                            res_table[i0, i1, i2, i3, i4, i5] = res
    return res_table


def analyze_gnn(res):
    embed_dim_list = [16, 32, 64]
    lr_list = [0.1, 0.01, 0.001]
    embedding_type_list = ["TransE", "TransR"]
    laplacian_type_list = ["symmetric", "random-walk"]
    aggregation_type_list = ["gcn", "graphsage", "lightgcn"]
    layers_list = [2, 4, 8]

    idx = np.unravel_index(np.argmax(res), res.shape)

    label_list = (
        embed_dim_list,
        lr_list,
        embedding_type_list,
        laplacian_type_list,
        aggregation_type_list,
        layers_list,
    )
    label_name = ("embed_dim", "lr", "embedding_type", "laplacian_type", "aggregation_type", "nlayers")
    label = get_label(idx, label_list, label_name)
    for dim in range(len(label_list)):
        plot_dim(res, dim, idx, label_list, "gnn", label_name[dim])
    return res[idx], label


def label2str(label: dict):
    s = []
    for k, v in label.items():
        s.append(f"{k}={v}")
    return "\n".join(s)


def analyze():
    base_res = get_baseline()
    base_best_recall, base_label = analyze_base(base_res)
    embed_res = get_embedding()
    embed_best_recall, embed_label = analyze_embedding(embed_res)
    gnn_res = get_gnn()
    gnn_best_recall, gnn_label = analyze_gnn(gnn_res)
    fig, ax = plt.subplots()
    x = ["base", "embedding", "gnn"]
    y = [base_best_recall, embed_best_recall, gnn_best_recall]
    label = [label2str(d) for d in [base_label, embed_label, gnn_label]]
    ax.scatter(x, y)
    ax.set_title("summary")
    ax.set_xlabel("methods")
    adjust_text([ax.text(i, j, f"{j:.4f}\n" + l) for i, j, l in zip(x, y, label)])
    fig.savefig("./res/summary.png")


if __name__ == "__main__":
    # run_baseline()
    # run_embedding()
    # run_gnn()
    analyze()
